import argparse
import time

import numpy as np
import pygame

from agents.dqn_agent import DQNAgent
from game_engine import WatermelonEnv
from training.train import train_dqn


def preprocess_observation(observation):
    """预处理观察，转换为(C, H, W)格式"""
    # 注意通道数变成了3*stack_frames
    return np.transpose(observation, (2, 0, 1)) / 255.0


def train(
    render_mode="rgb_array",
    render_every=100,
    episodes=5000,
    save_path="models",
    frame_skip=1,
    num_envs=1,
):
    """训练DQN智能体"""
    print("开始训练DQN智能体...")
    agent, rewards = train_dqn(
        episodes=episodes,
        render_every=render_every,
        save_path=save_path,
        render_mode=render_mode,
        frame_skip=frame_skip,
        num_envs=num_envs,
    )
    print("训练完成!")
    return agent


def test(model_path, render_mode="human", delay=0.01, frame_skip=1):
    """测试已训练的DQN智能体"""
    print(f"测试模型: {model_path}")

    # 创建环境
    env = WatermelonEnv(render_mode=render_mode, frame_skip=frame_skip)

    # 获取状态和动作空间信息
    observation, _ = env.reset()
    state_shape = (12, observation.shape[0], observation.shape[1])  # (C, H, W)
    n_actions = env.action_space.n

    # 创建并加载智能体
    agent = DQNAgent(state_shape, n_actions)
    agent.load(model_path)

    # 运行测试情节
    total_reward = 0
    episodes = 5

    for episode in range(episodes):
        observation, _ = env.reset()
        state = preprocess_observation(observation)
        episode_reward = 0

        while True:
            # 选择动作
            action = agent.select_action(state, training=False)

            # 执行动作
            observation, reward, terminated, truncated, info = env.step(action)
            episode_reward += reward

            # 获取下一个状态
            next_state = preprocess_observation(observation)

            # 更新状态
            state = next_state

            # 显示分数
            if render_mode == "human":
                pygame.display.set_caption(f"合成大西瓜 (AI) | 得分: {info['score']}")

            # 等待一小段时间，以便观察
            if delay > 0:
                time.sleep(delay)

            # 检查是否结束
            if terminated or truncated:
                total_reward += episode_reward
                print(
                    f"情节 {episode + 1} - 奖励: {episode_reward:.2f}, 得分: {info['score']}"
                )
                break

    env.close()
    print(f"平均奖励: {total_reward / episodes:.2f}")


def main():
    parser = argparse.ArgumentParser(
        description="Train or evaluate DQN for Watermelon Game"
    )
    parser.add_argument(
        "--mode",
        type=str,
        default="train",
        choices=["train", "test"],
        help="训练或测试模式",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="models/dqn_watermelon_final.pth",
        help="测试模式下的模型路径",
    )
    # 简化渲染控制参数
    parser.add_argument(
        "--render-mode",
        type=str,
        default="none",
        choices=["human", "rgb_array", "none"],
        help="渲染模式 (human: 窗口显示, rgb_array: 后台渲染, none: 不渲染)",
    )
    parser.add_argument(
        "--render-every",
        type=int,
        default=1,
        help="训练模式下，每隔多少个episode渲染一次",
    )
    parser.add_argument(
        "--episodes",
        type=int,
        default=500,
        help="训练的总episode数",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.01,
        help="测试模式下渲染每帧之间的延迟时间",
    )
    # 添加帧跳过参数
    parser.add_argument(
        "--frame-skip",
        type=int,
        default=10,
        help="每个动作重复的物理更新次数（帧跳过），提高训练效率",
    )
    parser.add_argument(
        "--num_envs",
        type=int,
        default=4,
        help="Number of parallel environments",
    )

    args = parser.parse_args()

    # 简化渲染模式处理
    render_mode = args.render_mode
    if render_mode == "none":
        render_mode = None
    elif render_mode is None:
        # 没有指定时使用默认值
        render_mode = "rgb_array" if args.mode == "train" else "human"

    if args.mode == "train":
        train(
            render_mode=render_mode,
            render_every=args.render_every,
            episodes=args.episodes,
            frame_skip=args.frame_skip,
            num_envs=args.num_envs,
        )
    else:
        test(
            args.model,
            render_mode=render_mode,
            delay=args.delay,
            frame_skip=args.frame_skip,
        )


if __name__ == "__main__":
    main()
